from flask import Flask, render_template, request, jsonify
import docker
import time
import threading
import pymysql
from datetime import datetime
import socket
import re

app = Flask(__name__)

# Docker客户端配置
client = docker.DockerClient(base_url='tcp://192.168.88.66:2375')

# 数据库配置 - 使用PyMySQL
db_config = {
    'host': '192.168.88.66',
    'user': 'root',
    'password': '202011',
    'database': 'ctf_platform',
    'charset': 'utf8mb4',
    'cursorclass': pymysql.cursors.DictCursor
}

# 全局端口池管理
PORT_POOL_MIN = 8001
PORT_POOL_MAX = 15000
port_lock = threading.Lock()
used_ports = set()

# 存储活动容器
active_containers = {}

# 靶场配置示例
TARGETS = {
    "cisp": {
        "image": "cisp:1.0",
        "ports": {
            "81/tcp": None,
            "82/tcp": None,
            "83/tcp": None,
            "84/tcp": None,
            "85/tcp": None
        },
        "key": "FLAG{Web_Sec_101}"
    },
    "syn": {
        "image": "syn:1.0",
        "ports": {"80/tcp": None},
        "key": "FLAG{Pwn_Master}"
    }
}


def get_db_connection():
    return pymysql.connect(**db_config)


def initialize_port_pool():
    """初始化端口池"""
    global used_ports

    # 从数据库获取已使用的端口
    conn = get_db_connection()
    try:
        with conn.cursor() as cursor:
            cursor.execute("SELECT ports FROM instances WHERE status='running'")
            for row in cursor.fetchall():
                ports_dict = eval(row['ports'])
                for port in ports_dict.values():
                    used_ports.add(port)
    finally:
        conn.close()

    # 检查当前活动容器的端口
    for container_info in active_containers.values():
        for port in container_info['ports'].values():
            used_ports.add(port)


def manage_container_timeout(container_id, timeout=3600):
    """容器超时管理线程"""
    time.sleep(timeout)
    if container_id in active_containers:
        stop_container(container_id)


def stop_container(container_id):
    """停止并移除容器"""
    if container_id in active_containers:
        try:
            container = client.containers.get(container_id)
            container.stop()
            container.remove()

            # 更新数据库
            conn = get_db_connection()
            try:
                with conn.cursor() as cursor:
                    cursor.execute(
                        "UPDATE instances SET status='stopped', end_time=%s WHERE container_id=%s",
                        (datetime.now(), container_id)
                    )
                conn.commit()
            finally:
                conn.close()

            # 释放端口
            with port_lock:
                ports_to_release = active_containers[container_id]["ports"].values()
                for port in ports_to_release:
                    if port in used_ports:
                        used_ports.remove(port)

            # 从活动容器中移除
            del active_containers[container_id]
            print(f"Container {container_id} stopped and removed")
        except docker.errors.NotFound:
            print(f"Container {container_id} not found")
        except docker.errors.APIError as e:
            print(f"Error stopping container {container_id}: {e}")


def recover_running_containers():
    """恢复运行中的容器状态"""
    conn = get_db_connection()
    try:
        with conn.cursor() as cursor:
            cursor.execute("SELECT * FROM instances WHERE status='running'")
            for row in cursor.fetchall():
                container_id = row['container_id']
                start_time = row['start_time'].timestamp()  # 转换为时间戳

                # 检查容器是否真的在运行
                try:
                    container = client.containers.get(container_id)
                    if container.status == 'running':
                        # 恢复容器信息
                        container_info = {
                            "id": container_id,
                            "name": row['target_name'],
                            "ports": eval(row['ports']),
                            "key": TARGETS.get(row['target_name'], {}).get("key", ""),
                            "start_time": start_time  # 存储原始时间戳
                        }
                        active_containers[container_id] = container_info

                        # 计算剩余时间（秒）
                        elapsed = time.time() - start_time
                        remaining_time = max(0, 3600 - elapsed)

                        if remaining_time > 0:
                            threading.Thread(
                                target=manage_container_timeout,
                                args=(container_id, remaining_time)
                            ).start()
                except docker.errors.NotFound:
                    # 容器不存在，更新数据库状态
                    try:
                        with conn.cursor() as update_cursor:
                            update_cursor.execute(
                                "UPDATE instances SET status='stopped' WHERE container_id=%s",
                                (container_id,)
                            )
                        conn.commit()
                    except Exception as e:
                        print(f"更新容器状态失败: {e}")
    finally:
        conn.close()


def is_port_available(port):
    """检查端口是否可用（未被系统占用）"""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        try:
            s.bind(('0.0.0.0', port))
            return True
        except OSError:
            return False


def find_available_port():
    """查找可用端口"""
    global used_ports

    with port_lock:
        # 尝试从8001-15000范围内查找
        for port in range(PORT_POOL_MIN, PORT_POOL_MAX + 1):
            if port not in used_ports and is_port_available(port):
                used_ports.add(port)
                return port

        # 如果常用端口范围找不到，尝试其他范围
        for port in range(15001, 20000):
            if port not in used_ports and is_port_available(port):
                used_ports.add(port)
                return port

        raise Exception("没有可用端口")


# 添加此函数（放在 TARGETS 定义之后）
def get_target_display_name(target_name):
    display_names = {
        "cisp": "CISP靶场",
        "syn": "综合靶场"
    }
    return display_names.get(target_name, target_name)


# 初始化时恢复运行中的容器
initialize_port_pool()
recover_running_containers()


@app.route('/')
def index():
    """主页显示可用靶场和运行中的容器"""
    # 获取所有运行中的容器信息
    running_containers = []
    for container_id, info in active_containers.items():
        running_containers.append({
            "id": container_id,
            "name": info["name"],
            "ports": info["ports"]
        })

    return render_template('index.html',
                           targets=TARGETS.keys(),
                           running_containers=running_containers,
                           get_target_display_name=get_target_display_name)


@app.route('/start_target', methods=['POST'])
def start_target():
    """启动靶场容器"""
    target_name = request.form['target']
    target_config = TARGETS.get(target_name)

    if not target_config:
        return jsonify({"error": "Invalid target"}), 400

    # 分配端口
    allocated_ports = {}
    for container_port, _ in target_config["ports"].items():
        try:
            host_port = find_available_port()
            allocated_ports[container_port] = host_port
        except Exception as e:
            return jsonify({"error": f"端口分配失败: {str(e)}"}), 500

    try:
        container = client.containers.run(
            target_config["image"],
            detach=True,
            ports=allocated_ports
        )
        # 存储容器信息
        container_info = {
            "id": container.id,
            "name": target_name,
            "ports": allocated_ports,
            "key": target_config["key"],
            "start_time": time.time()
        }
        active_containers[container.id] = container_info

        # 启动超时线程
        threading.Thread(target=manage_container_timeout, args=(container.id,)).start()

        # 保存到数据库
        conn = get_db_connection()
        try:
            with conn.cursor() as cursor:
                sql = """
                    INSERT INTO instances (container_id, target_name, ports, start_time, status) 
                    VALUES (%s, %s, %s, %s, 'running')
                    """
                cursor.execute(sql, (
                    container.id,
                    target_name,
                    str(allocated_ports),
                    datetime.now()
                ))
            conn.commit()
        finally:
            conn.close()

        return jsonify({
            "id": container.id,
            "ports": allocated_ports,
            "message": f"Target {target_name} started"
        })

    except docker.errors.ImageNotFound:
        # 释放已分配端口
        with port_lock:
            for port in allocated_ports.values():
                if port in used_ports:
                    used_ports.remove(port)
        return jsonify({"error": "Docker镜像未找到"}), 500
    except docker.errors.APIError as e:
        # 释放已分配端口
        with port_lock:
            for port in allocated_ports.values():
                if port in used_ports:
                    used_ports.remove(port)
        return jsonify({"error": f"Docker API错误: {str(e)}"}), 500
    except Exception as e:
        # 释放已分配端口
        with port_lock:
            for port in allocated_ports.values():
                if port in used_ports:
                    used_ports.remove(port)
        return jsonify({"error": f"未知错误: {str(e)}"}), 500


@app.route('/stop_target', methods=['POST'])
def stop_target():
    """停止靶场容器"""
    container_id = request.form['container_id']
    if container_id in active_containers:
        stop_container(container_id)
        return jsonify({"message": f"Container {container_id} stopped"})
    else:
        return jsonify({"error": "Container not found or already stopped"}), 404


@app.route('/submit_key', methods=['POST'])
def submit_key():
    """验证用户提交的Key"""
    container_id = request.form['container_id']
    user_key = request.form['key']

    if container_id not in active_containers:
        return jsonify({"valid": False, "message": "Container not found or already stopped"})

    target_key = active_containers[container_id]["key"]
    is_valid = (user_key.strip() == target_key)

    if is_valid:
        # 更新数据库记录
        conn = get_db_connection()
        try:
            with conn.cursor() as cursor:
                cursor.execute(
                    "UPDATE instances SET completed=1 WHERE container_id=%s",
                    (container_id,)
                )
            conn.commit()
        finally:
            conn.close()

        # 停止容器
        stop_container(container_id)

    return jsonify({"valid": is_valid})


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)
